# coding='utf-8'
import os
import sys
import numpy as np
import time
import datetime
import json
import importlib
import logging
import shutil
import cv2
import random
from tqdm import tqdm
import matplotlib

matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.ticker import NullLocator

import torch
import torch.nn as nn

MY_DIRNAME = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(MY_DIRNAME, '..'))
from nets.model_main import ModelMain
from nets.yolo_loss import YOLOLoss
from common.utils import non_max_suppression, bbox_iou

cmap = plt.get_cmap('tab20b')
colors = [cmap(i) for i in np.linspace(0, 1, 20)]


def test(config):
    is_training = False
    # Load and initialize network
    net = ModelMain(config, is_training=is_training)
    net.train(is_training)

    # Set data parallel
    net = nn.DataParallel(net)
    net = net.cuda()

    # Restore pretrain model
    if config["pretrain_snapshot"]:
        logging.info("load checkpoint from {}".format(config["pretrain_snapshot"]))
        state_dict = torch.load(config["pretrain_snapshot"])
        net.load_state_dict(state_dict)
    else:
        raise Exception("missing pretrain_snapshot!!!")

    # YOLO loss with 3 scales
    yolo_losses = []
    for i in range(3):
        yolo_losses.append(YOLOLoss(config["yolo"]["anchors"][i],
                                    config["yolo"]["classes"], (config["img_w"], config["img_h"])))

    # prepare images path
    images_name = os.listdir(config["images_path"])
    images_path = [os.path.join(config["images_path"], name) for name in images_name]
    if len(images_path) == 0:
        raise Exception("no image found in {}".format(config["images_path"]))
    if not os.path.isdir("/home/iotcam/Pictures/video/seg_pics6/"):
        os.makedirs("/home/iotcam/Pictures/video/seg_pics6/")
    index = 0
    # Start inference
    for path in tqdm(images_path):
        image = cv2.imread(path, cv2.IMREAD_COLOR)
        image_origin = image  # keep for save result
        if image is None:
            logging.error("read path error: {}. skip it.".format(path))
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = cv2.resize(image, (config["img_w"], config["img_h"]),
                           interpolation=cv2.INTER_LINEAR)
        image = image.astype(np.float32)
        image /= 255.0
        image = np.transpose(image, (2, 0, 1))
        image = image.astype(np.float32)
        image = np.asarray(image)
        image = torch.from_numpy(image).unsqueeze(0).cuda()
        # inference
        with torch.no_grad():
            outputs = net(image)
            output_list = []
            for i in range(3):
                output_list.append(yolo_losses[i](outputs[i]))
            output = torch.cat(output_list, 1)
            batch_detections = non_max_suppression(output, config["yolo"]["classes"],
                                                   conf_thres=config["confidence_threshold"],
                                                   nms_thres=0.25)

        for detections in batch_detections:

            if detections is not None:
                for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
                    if int(cls_pred) < 2:
                        # Rescale coordinates to original dimensions
                        ori_h, ori_w = image_origin.shape[:2]
                        pre_h, pre_w = config["img_h"], config["img_w"]
                        box_h = ((y2 - y1) / pre_h) * ori_h
                        box_w = ((x2 - x1) / pre_w) * ori_w
                        y1 = (y1 / pre_h) * ori_h
                        x1 = (x1 / pre_w) * ori_w
                        y2 = y1 + box_h
                        x2 = x1 + box_w
                        x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
                        if x1 < 0:
                            x1 = 0
                        if y1 < 0:
                            y1 = 0
                        if x2 > 1920:
                            x2 = 1920
                        if y2 > 1080:
                            y2 = 1080
                        image_cls = image_origin[y1:y2, x1:x2]  # # img[height, width]
                        img_name = '0923_6_%06d.jpg' % index
                        save_path = os.path.join("/home/iotcam/Pictures/video/seg_pics6/", img_name)
                        cv2.imwrite(save_path, image_cls)
                        index += 1


def main():
    logging.basicConfig(level=logging.WARN,
                        format="[%(asctime)s %(filename)s] %(message)s")

    if len(sys.argv) != 2:
        logging.error("Usage: python test_images.py params.py")
        sys.exit()
    params_path = sys.argv[1]
    if not os.path.isfile(params_path):
        logging.error("no params file found! path: {}".format(params_path))
        sys.exit()
    config = importlib.import_module(params_path[:-3]).TRAINING_PARAMS
    config["batch_size"] *= len(config["parallels"])

    # Start training
    os.environ["CUDA_VISIBLE_DEVICES"] = ','.join(map(str, config["parallels"]))
    test(config)


if __name__ == "__main__":
    main()
